"""
Excel导出模块
负责将数据库结构导出为Excel文件，支持AI增强标识和样式设置
"""

import logging
from typing import Dict, List, Any, Optional, Tuple
import pandas as pd
from openpyxl import Workbook
from openpyxl.styles import Font, Fill, PatternFill, Alignment, Border, Side
from openpyxl.utils.dataframe import dataframe_to_rows
from openpyxl.utils import get_column_letter
from openpyxl.comments import Comment
import os
from datetime import datetime

from src.database.schema_extractor import DatabaseSchema, TableInfo, ColumnInfo


class ExcelStyleConfig:
    """Excel样式配置"""
    
    # 颜色定义
    HEADER_COLOR = "366092"  # 深蓝色
    AI_ENHANCED_COLOR = "E6F3FF"  # 浅蓝色背景
    PRIMARY_KEY_COLOR = "FFE6E6"  # 浅红色
    FOREIGN_KEY_COLOR = "E6FFE6"  # 浅绿色
    WARNING_COLOR = "FFF2CC"  # 浅黄色
    
    # 字体定义
    HEADER_FONT = Font(name='微软雅黑', size=11, bold=True, color="FFFFFF")
    NORMAL_FONT = Font(name='微软雅黑', size=10)
    AI_FONT = Font(name='微软雅黑', size=10, italic=True)
    
    # 填充样式
    HEADER_FILL = PatternFill(start_color=HEADER_COLOR, end_color=HEADER_COLOR, fill_type="solid")
    AI_FILL = PatternFill(start_color=AI_ENHANCED_COLOR, end_color=AI_ENHANCED_COLOR, fill_type="solid")
    PK_FILL = PatternFill(start_color=PRIMARY_KEY_COLOR, end_color=PRIMARY_KEY_COLOR, fill_type="solid")
    FK_FILL = PatternFill(start_color=FOREIGN_KEY_COLOR, end_color=FOREIGN_KEY_COLOR, fill_type="solid")
    
    # 对齐方式
    CENTER_ALIGNMENT = Alignment(horizontal='center', vertical='center')
    LEFT_ALIGNMENT = Alignment(horizontal='left', vertical='center', wrap_text=True)
    
    # 边框
    THIN_BORDER = Border(
        left=Side(style='thin'),
        right=Side(style='thin'),
        top=Side(style='thin'),
        bottom=Side(style='thin')
    )


class ExcelExporter:
    """Excel导出器"""
    
    def __init__(self):
        self.logger = logging.getLogger(__name__)
        self.style_config = ExcelStyleConfig()
    
    def export_database_schema(self, schema: DatabaseSchema, 
                             output_path: str,
                             include_summary: bool = True,
                             include_relationships: bool = True,
                             include_data_info: bool = True,
                             group_by_schema: bool = True,
                             ai_enhanced: bool = False) -> bool:
        """
        导出数据库结构到Excel
        
        Args:
            schema: 数据库结构对象
            output_path: 输出文件路径
            include_summary: 是否包含汇总信息
            include_relationships: 是否包含关系分析
            include_data_info: 是否包含数据信息（行数、大小等）
            group_by_schema: 是否按schema分组
            ai_enhanced: 是否启用了AI增强（控制是否显示AI标识）
        """
        try:
            self.logger.info(f"开始导出数据库结构到Excel: {output_path}")
            
            # 创建工作簿
            wb = Workbook()
            
            # 删除默认工作表
            wb.remove(wb.active)
            
            # 1. 创建汇总工作表
            if include_summary:
                self._create_summary_sheet(wb, schema, ai_enhanced)
            
            # 2. 创建表结构工作表
            if group_by_schema:
                # 按schema分组创建工作表
                for schema_name, tables in schema.schemas.items():
                    if tables:
                        self._create_schema_tables_sheet(wb, schema_name, tables, 
                                                       include_relationships, include_data_info, ai_enhanced)
            else:
                # 所有表放在一个工作表中
                all_tables = []
                for tables in schema.schemas.values():
                    all_tables.extend(tables)
                if all_tables:
                    self._create_schema_tables_sheet(wb, "所有表结构", all_tables, 
                                                   include_relationships, include_data_info, ai_enhanced)
            
            # 3. 创建关系分析工作表
            if include_relationships:
                self._create_relationships_sheet(wb, schema, ai_enhanced)
            
            # 4. 创建元数据工作表
            self._create_metadata_sheet(wb, schema)
            
            # 保存文件
            wb.save(output_path)
            
            self.logger.info(f"Excel导出成功: {output_path}")
            return True
            
        except Exception as e:
            self.logger.error(f"Excel导出失败: {str(e)}")
            return False
    
    def _create_summary_sheet(self, wb: Workbook, schema: DatabaseSchema, ai_enhanced: bool = False):
        """创建汇总工作表"""
        ws = wb.create_sheet("数据库汇总")
        
        # 准备汇总数据
        summary_data = []
        
        # 基本信息
        summary_data.append(["数据库名称", schema.database_name])
        summary_data.append(["数据库类型", schema.database_type.value])
        summary_data.append(["配置名称", schema.config_name])
        summary_data.append(["提取时间", schema.extraction_time or "未知"])
        summary_data.append(["总表数", schema.total_tables])
        summary_data.append(["总字段数", schema.total_columns])
        summary_data.append(["", ""])  # 空行
        
        # Schema统计
        if ai_enhanced:
            summary_data.append(["Schema名称", "表数量", "字段数量", "AI增强表数", "AI增强字段数"])
            
            for schema_name, tables in schema.schemas.items():
                table_count = len(tables)
                column_count = sum(len(table.columns) for table in tables)
                
                # 统计AI增强情况
                ai_table_count = sum(1 for table in tables 
                                   if table.table_comment and self._is_ai_generated(table.table_comment))
                ai_column_count = sum(1 for table in tables 
                                    for column in table.columns 
                                    if column.column_comment and self._is_ai_generated(column.column_comment))
                
                summary_data.append([schema_name, table_count, column_count, ai_table_count, ai_column_count])
        else:
            summary_data.append(["Schema名称", "表数量", "字段数量"])
            
            for schema_name, tables in schema.schemas.items():
                table_count = len(tables)
                column_count = sum(len(table.columns) for table in tables)
                summary_data.append([schema_name, table_count, column_count])
        
        # 写入数据
        for row_idx, row_data in enumerate(summary_data, 1):
            for col_idx, value in enumerate(row_data, 1):
                cell = ws.cell(row=row_idx, column=col_idx, value=value)
                
                # 设置样式
                # 检查是否为表头行（包含"Schema名称"的行）
                is_header_row = any("Schema名称" in str(cell_value) for cell_value in row_data)
                if is_header_row:
                    cell.font = self.style_config.HEADER_FONT
                    cell.fill = self.style_config.HEADER_FILL
                    cell.alignment = self.style_config.CENTER_ALIGNMENT
                else:
                    cell.font = self.style_config.NORMAL_FONT
                    cell.alignment = self.style_config.LEFT_ALIGNMENT
                
                cell.border = self.style_config.THIN_BORDER
        
        # 调整列宽
        for column in ws.columns:
            max_length = 0
            column_letter = get_column_letter(column[0].column)
            for cell in column:
                try:
                    if len(str(cell.value)) > max_length:
                        max_length = len(str(cell.value))
                except:
                    pass
            adjusted_width = min(max_length + 2, 50)
            ws.column_dimensions[column_letter].width = adjusted_width
    
    def _create_schema_tables_sheet(self, wb: Workbook, sheet_name: str, tables: List[TableInfo],
                                  include_relationships: bool, include_data_info: bool, ai_enhanced: bool = False):
        """创建表结构工作表"""
        ws = wb.create_sheet(sheet_name[:31])  # Excel工作表名称限制31字符
        
        # 定义列标题
        headers = []
        if ai_enhanced and include_relationships:
            headers.append("表关系分析")
        headers.extend(["表名", "表注释", "字段名", "字段注释", "数据类型", "是否为空", "默认值", "主键", "外键", "自增"])
        
        if include_data_info:
            headers.extend(["行数", "表大小"])
        
        # 写入表头
        for col_idx, header in enumerate(headers, 1):
            cell = ws.cell(row=1, column=col_idx, value=header)
            cell.font = self.style_config.HEADER_FONT
            cell.fill = self.style_config.HEADER_FILL
            cell.alignment = self.style_config.CENTER_ALIGNMENT
            cell.border = self.style_config.THIN_BORDER
        
        # 写入数据
        current_row = 2
        
        for table in tables:
            # 第一行写入表信息
            first_column = table.columns[0] if table.columns else None
            
            row_data = []
            
            # 根据AI增强状态决定是否包含表关系分析
            if ai_enhanced and include_relationships:
                relationship_analysis = self._extract_relationship_analysis(table.table_comment)
                row_data.append(relationship_analysis)
            
            # 基本表和字段信息
            table_comment_clean = self._clean_table_comment(table.table_comment)
            column_comment_clean = self._clean_display_text(first_column.column_comment) if first_column else ""
            
            row_data.extend([
                table.table_name,       # 表名
                self._clean_display_text(table_comment_clean),  # 表注释（去除关系分析部分和AI标记）
                first_column.column_name if first_column else "",  # 字段名
                column_comment_clean,  # 字段注释（去除AI标记）
                first_column.data_type if first_column else "",  # 数据类型
                "是" if first_column and first_column.nullable else "否",  # 是否为空
                first_column.default_value if first_column else "",  # 默认值
                "是" if first_column and first_column.is_primary_key else "否",  # 主键
                "是" if first_column and first_column.is_foreign_key else "否",  # 外键
                "是" if first_column and first_column.auto_increment else "否"   # 自增
            ])
            
            if include_data_info:
                row_data.extend([
                    table.row_count or "",  # 行数
                    table.table_size or ""  # 表大小
                ])
            
            # 写入第一行
            self._write_table_row(ws, current_row, row_data, table, first_column, True, ai_enhanced, include_relationships)
            current_row += 1
            
            # 写入其余字段行
            for column in table.columns[1:]:
                row_data = []
                
                # 根据AI增强状态决定是否包含表关系分析（空白）
                if ai_enhanced and include_relationships:
                    row_data.append("")  # 表关系分析（只在第一行显示）
                
                # 基本信息
                row_data.extend([
                    "",  # 表名（只在第一行显示）
                    "",  # 表注释（只在第一行显示）
                    column.column_name,    # 字段名
                    self._clean_display_text(column.column_comment), # 字段注释（去除AI标记）
                    column.data_type,      # 数据类型
                    "是" if column.nullable else "否",  # 是否为空
                    column.default_value or "",  # 默认值
                    "是" if column.is_primary_key else "否",  # 主键
                    "是" if column.is_foreign_key else "否",  # 外键
                    "是" if column.auto_increment else "否"   # 自增
                ])
                
                if include_data_info:
                    row_data.extend(["", ""])  # 数据信息只在第一行显示
                
                self._write_table_row(ws, current_row, row_data, table, column, False, ai_enhanced, include_relationships)
                current_row += 1
            
            # 添加空行分隔不同表
            current_row += 1
        
        # 调整列宽
        self._adjust_column_widths(ws)
        
        # 冻结表头
        ws.freeze_panes = "A2"
    
    def _write_table_row(self, ws, row_idx: int, row_data: List[str], 
                        table: TableInfo, column: Optional[ColumnInfo], is_first_row: bool,
                        ai_enhanced: bool = False, include_relationships: bool = False):
        """写入表行数据并设置样式"""
        for col_idx, value in enumerate(row_data, 1):
            cell = ws.cell(row=row_idx, column=col_idx, value=value)
            
            # 设置基本样式
            cell.font = self.style_config.NORMAL_FONT
            cell.alignment = self.style_config.LEFT_ALIGNMENT
            cell.border = self.style_config.THIN_BORDER
            
            # 根据列类型设置特殊样式（仅在AI增强模式下）
            if ai_enhanced:
                # 动态计算列索引偏移
                offset = 0
                if ai_enhanced and include_relationships:
                    # 表关系分析列
                    if col_idx == 1:
                        if value and self._is_ai_generated(value):
                            cell.fill = self.style_config.AI_FILL
                            cell.font = self.style_config.AI_FONT
                            cell.comment = Comment("AI生成的关系分析", "系统")
                    offset = 1
                
                # 表注释列
                table_comment_col = 2 + offset
                if col_idx == table_comment_col:
                    if value and self._is_ai_generated(table.table_comment):
                        cell.fill = self.style_config.AI_FILL
                        cell.font = self.style_config.AI_FONT
                        cell.comment = Comment("AI生成的表注释", "系统")
                
                # 字段注释列
                column_comment_col = 4 + offset
                if col_idx == column_comment_col and column:
                    if value and self._is_ai_generated(column.column_comment):
                        cell.fill = self.style_config.AI_FILL
                        cell.font = self.style_config.AI_FONT
                        cell.comment = Comment("AI生成的字段注释", "系统")
            
            # 主键和外键列样式（考虑偏移）
            offset = 1 if (ai_enhanced and include_relationships) else 0
            primary_key_col = 8 + offset  # 主键列
            foreign_key_col = 9 + offset  # 外键列
            
            if col_idx == primary_key_col and column and column.is_primary_key:
                cell.fill = self.style_config.PK_FILL
            elif col_idx == foreign_key_col and column and column.is_foreign_key:
                cell.fill = self.style_config.FK_FILL
    
    def _create_relationships_sheet(self, wb: Workbook, schema: DatabaseSchema, ai_enhanced: bool = False):
        """创建关系分析工作表"""
        ws = wb.create_sheet("表关系分析")
        
        # 表头
        headers = ["Schema", "表名", "表注释", "关系分析", "外键关系", "被引用关系"]
        
        for col_idx, header in enumerate(headers, 1):
            cell = ws.cell(row=1, column=col_idx, value=header)
            cell.font = self.style_config.HEADER_FONT
            cell.fill = self.style_config.HEADER_FILL
            cell.alignment = self.style_config.CENTER_ALIGNMENT
            cell.border = self.style_config.THIN_BORDER
        
        # 数据行
        current_row = 2
        
        for schema_name, tables in schema.schemas.items():
            for table in tables:
                # 外键关系
                fk_relations = []
                for fk in table.foreign_keys:
                    fk_relations.append(f"{fk.column_name} -> {fk.referenced_table}.{fk.referenced_column}")
                
                # 被引用关系（需要查找其他表）
                referenced_by = []
                for other_table in tables:
                    if other_table.table_name == table.table_name:
                        continue
                    for fk in other_table.foreign_keys:
                        if fk.referenced_table == table.table_name:
                            referenced_by.append(f"{other_table.table_name}.{fk.column_name}")
                
                row_data = [
                    schema_name,
                    table.table_name,
                    self._clean_display_text(self._clean_table_comment(table.table_comment)),
                    self._clean_display_text(self._extract_relationship_analysis(table.table_comment)),
                    "; ".join(fk_relations),
                    "; ".join(referenced_by)
                ]
                
                for col_idx, value in enumerate(row_data, 1):
                    cell = ws.cell(row=current_row, column=col_idx, value=value)
                    cell.font = self.style_config.NORMAL_FONT
                    cell.alignment = self.style_config.LEFT_ALIGNMENT
                    cell.border = self.style_config.THIN_BORDER
                    
                    # AI生成内容标识（仅在AI增强模式下）
                    if ai_enhanced and col_idx == 4 and value and self._is_ai_generated(table.table_comment):
                        cell.fill = self.style_config.AI_FILL
                        cell.font = self.style_config.AI_FONT
                
                current_row += 1
        
        # 调整列宽
        self._adjust_column_widths(ws)
    
    def _create_metadata_sheet(self, wb: Workbook, schema: DatabaseSchema):
        """创建元数据工作表"""
        ws = wb.create_sheet("导出信息")
        
        metadata = [
            ["导出时间", datetime.now().strftime("%Y-%m-%d %H:%M:%S")],
            ["数据库名称", schema.database_name],
            ["数据库类型", schema.database_type.value],
            ["配置名称", schema.config_name],
            ["结构提取时间", schema.extraction_time or "未知"],
            ["总表数", schema.total_tables],
            ["总字段数", schema.total_columns],
            ["", ""],
            ["AI增强标识说明", ""],
            ["浅蓝色背景", "AI生成的内容"],
            ["浅红色背景", "主键字段"],
            ["浅绿色背景", "外键字段"],
            ["", ""],
            ["导出工具", "数据预处理系统 v1.0"],
            ["导出格式", "Excel (.xlsx)"],
        ]
        
        for row_idx, (key, value) in enumerate(metadata, 1):
            ws.cell(row=row_idx, column=1, value=key).font = self.style_config.NORMAL_FONT
            ws.cell(row=row_idx, column=2, value=value).font = self.style_config.NORMAL_FONT
            
            # 特殊样式
            if key == "浅蓝色背景":
                ws.cell(row=row_idx, column=2).fill = self.style_config.AI_FILL
            elif key == "浅红色背景":
                ws.cell(row=row_idx, column=2).fill = self.style_config.PK_FILL
            elif key == "浅绿色背景":
                ws.cell(row=row_idx, column=2).fill = self.style_config.FK_FILL
        
        ws.column_dimensions['A'].width = 20
        ws.column_dimensions['B'].width = 40
    
    def _adjust_column_widths(self, ws):
        """调整列宽"""
        for column in ws.columns:
            max_length = 0
            column_letter = get_column_letter(column[0].column)
            
            for cell in column:
                try:
                    if cell.value:
                        max_length = max(max_length, len(str(cell.value)))
                except:
                    pass
            
            # 设置合适的列宽
            if column_letter == 'A':  # 表关系分析列
                adjusted_width = min(max_length + 2, 30)
            elif column_letter in ['B', 'C']:  # 表名、表注释列
                adjusted_width = min(max_length + 2, 25)
            elif column_letter == 'E':  # 字段注释列
                adjusted_width = min(max_length + 2, 20)
            else:
                adjusted_width = min(max_length + 2, 15)
            
            ws.column_dimensions[column_letter].width = max(adjusted_width, 8)
    
    def _is_ai_generated(self, text: str) -> bool:
        """判断文本是否为AI生成（基于明确的AI生成标记）"""
        if not text:
            return False
        
        # 检查是否包含AI生成标记
        return "[AI生成]" in text or "关系分析:" in text
    
    def _clean_display_text(self, text: str) -> str:
        """清理显示文本，移除AI生成标记"""
        if not text:
            return ""
        
        # 移除AI生成标记
        return text.replace("[AI生成]", "").strip()
    
    def _extract_relationship_analysis(self, table_comment: str) -> str:
        """从表注释中提取关系分析部分"""
        if not table_comment:
            return ""
        
        # 查找关系分析标记
        if "关系分析:" in table_comment:
            parts = table_comment.split("关系分析:")
            if len(parts) > 1:
                analysis = parts[1].strip()
                # 清理AI生成标记
                return self._clean_display_text(analysis)
        
        return ""
    
    def _clean_table_comment(self, table_comment: str) -> str:
        """清理表注释，移除关系分析部分"""
        if not table_comment:
            return ""
        
        # 移除关系分析部分
        if "关系分析:" in table_comment:
            parts = table_comment.split("关系分析:")
            return parts[0].strip().rstrip("|").strip()
        
        return table_comment
    
    def export_tables_to_dataframe(self, tables: List[TableInfo]) -> pd.DataFrame:
        """将表结构导出为DataFrame"""
        data = []
        
        for table in tables:
            for column in table.columns:
                row = {
                    "表名": table.table_name,
                    "表注释": self._clean_display_text(table.table_comment),
                    "字段名": column.column_name,
                    "字段注释": self._clean_display_text(column.column_comment),
                    "数据类型": column.data_type,
                    "是否为空": "是" if column.nullable else "否",
                    "默认值": column.default_value,
                    "主键": "是" if column.is_primary_key else "否",
                    "外键": "是" if column.is_foreign_key else "否",
                    "自增": "是" if column.auto_increment else "否"
                }
                data.append(row)
        
        return pd.DataFrame(data)


# 全局Excel导出器实例
excel_exporter = ExcelExporter()